from threading import Lock

from typing import TypeVar, Generic, List, Callable

from common.timeout import timeout

T = TypeVar('T')


class ObjectPool(Generic[T]):
    def __init__(self, object_factory: Callable[[], T], max_objects=10):
        self.object_factory = object_factory
        self.max_objects: int = max_objects
        self.available_objects: List[T] = []
        self.used_objects: List[T] = []
        self._lock = Lock()

    def __enter__(self):
        self._lock.acquire()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._lock.release()

    @timeout(3)
    def __get(self):
        obj: T = None
        if self.available_objects:
            obj = self.available_objects.pop()
        elif len(self.used_objects) < self.max_objects:
            obj = self.object_factory()
        else:
            while True:
                if self.available_objects:
                    obj = self.available_objects.pop()
                    break
        # 加入到已使用列表中
        self.used_objects.append(obj)
        return obj

    def acquire(self) -> T:
        if self._lock.locked():
            return self.__get()
        else:
            with self._lock:
                return self.__get()

    def __re(self, obj):
        if obj in self.used_objects:
            self.available_objects.append(obj)
            self.used_objects.remove(obj)

    def release(self, obj: T):
        if self._lock.locked():
            return self.__re(obj)
        else:
            with self._lock:
                return self.__re(obj)

    def size(self):
        return len(self.available_objects)
